// SPDX-License-Identifier: GPL-2.0

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/err.h>
#include <linux/slab.h>
#include <linux/fs.h>
#include <linux/vmalloc.h>

#include "dim_elf.h"

int dim_elf_get_ehdr(struct file *elf_file, struct elfhdr *ehdr)
{
	loff_t pos = 0;
	ssize_t size;

	if (elf_file == NULL || ehdr == NULL)
		return -EINVAL;

	size = kernel_read(elf_file, ehdr, sizeof(struct elfhdr), &pos);
	if (size != sizeof(struct elfhdr))
		return size < 0 ? (int)size : -EIO;

	/* check elf header valid */
	if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG) != 0 ||
	    ehdr->e_ident[EI_CLASS] != ELF_CLASS ||
#ifdef __LITTLE_ENDIAN__
	    ehdr->e_ident[EI_DATA]  == ELFDATA2LSB ||
#else
	    ehdr->e_ident[EI_DATA]  == ELFDATA2MSB ||
#endif
	    (ehdr->e_type != ET_EXEC && ehdr->e_type != ET_DYN))
		return -ENOEXEC;

	return 0;
}

int dim_elf_get_phdrs(struct file *elf_file,
		      struct elfhdr *ehdr,
		      struct elf_phdr **phdrs,
		      size_t *num)
{
	struct elf_phdr *elf_phdata = NULL;
	size_t phdr_size;
	ssize_t read_size;

	if (elf_file == NULL || ehdr == NULL || phdrs == NULL || num == NULL)
		return -EINVAL;

	if (ehdr->e_phentsize != sizeof(struct elf_phdr) ||
	    ehdr->e_phnum < 1 ||
	    ehdr->e_phnum > 65536U / sizeof(struct elf_phdr))
		return -ENOEXEC;

	phdr_size = sizeof(struct elf_phdr) * ehdr->e_phnum;
	elf_phdata = kmalloc(phdr_size, GFP_KERNEL);
	if (elf_phdata == NULL)
		return -ENOMEM;

	read_size = kernel_read(elf_file,
#ifdef __arm__
		elf_phdata, phdr_size, (loff_t *)&ehdr->e_phoff);
#else
		elf_phdata, phdr_size, &ehdr->e_phoff);
#endif
	if (read_size != phdr_size) {
		kfree(elf_phdata);
		return read_size < 0 ? (int)read_size : -EIO;
	}

	*phdrs = elf_phdata;
	*num = ehdr->e_phnum;
	return 0;
}

int elf_get_section_by_name(struct file *elf_file,
			    struct elfhdr *ehdr,
			    const char *name,
			    struct elf_shdr *shdr)
{
	int i, ret;
	ssize_t size, name_len, str_size;
	struct elf_shdr *sh_table = NULL;
	char *sh_str = NULL;
	loff_t pos;

	if (elf_file == NULL || ehdr == NULL || name == NULL || shdr == NULL)
		return -EINVAL;

	if (ehdr->e_shentsize != sizeof(struct elf_shdr))
		return -EBADF;

	sh_table = kmalloc(ehdr->e_shentsize, GFP_KERNEL);
	if (sh_table == NULL)
		return -ENOMEM;

	pos = ehdr->e_shoff + ehdr->e_shentsize * ehdr->e_shstrndx;
	size = kernel_read(elf_file, sh_table, ehdr->e_shentsize, &pos);
	if (size != ehdr->e_shentsize) {
		kfree(sh_table);
		return size < 0 ? (int)size : -EBADF;
	}

	str_size = sh_table->sh_size;
	if (str_size > i_size_read(file_inode(elf_file))) {
		kfree(sh_table);
		return -EBADF;
	}

	sh_str = vmalloc(str_size);
	if (sh_str == NULL) {
		kfree(sh_table);
		return -ENOMEM;
	}

	pos = sh_table->sh_offset;
	size = kernel_read(elf_file, sh_str, sh_table->sh_size, &pos);
	if (size != sh_table->sh_size) {
		kfree(sh_table);
		vfree(sh_str);
		return size < 0 ? (int)size : -EBADF;
	}

	ret = -ENOENT;
	pos = ehdr->e_shoff;
	name_len = strlen(name);
	for (i = 0; i < ehdr->e_shnum; i++) {
		size = kernel_read(elf_file, sh_table, ehdr->e_shentsize, &pos);
		if (size != ehdr->e_shentsize) {
			ret = size < 0 ? (int)size : -EBADF;
			break;
		}

		if (sh_table->sh_name + name_len < sh_table->sh_name ||
		    sh_table->sh_name + name_len >= str_size) {
			ret = -EBADF;
			break;
		}

		if (strcmp(name, sh_str + sh_table->sh_name) == 0) {
			memcpy(shdr, sh_table, sizeof(struct elf_shdr));
			ret = 0;
			break;
		}
	}

	kfree(sh_table);
	vfree(sh_str);
	return ret;
}
